# -*- coding: utf-8 -*-

import os
import sys
from PyQt5 import QtCore, QtGui
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QApplication, QDialog
from PyQt5.QtGui import QImage, QPixmap, QPainter, QPen
from od.check.od_check_ui import Ui_od_check
from od.check.od_check_backend import ODCheckBackend


class MatchViewWindow(QDialog, Ui_od_check):
    def __init__(self, folder):
        super(MatchViewWindow, self).__init__()
        self.setupUi(self)

        self.pushButton_previous_pair.clicked.connect(self.previous_pair_button)
        self.pushButton_next_pair.clicked.connect(self.next_pair_button)
        self.pushButton_gallery_previous.clicked.connect(self.previous_gallery_button)
        self.pushButton_gallery_next.clicked.connect(self.next_gallery_button)
        self.pushButton_query_previous.clicked.connect(self.previous_query_button)
        self.pushButton_query_next.clicked.connect(self.next_query_button)
        self.radioButton_same_person.clicked.connect(self.same_person_radio_button)
        self.radioButton_different_person.clicked.connect(self.different_person_radio_button)
        self.radioButton_uncertain.clicked.connect(self.uncertain_radio_button)
        self.backend = ODCheckBackend(folder)

        self.label_gallery.setScaledContents(True)  # 图像铺满Label
        self.label_query.setScaledContents(True)
        self.label_result_image.setScaledContents(True)
        self.show_image()

    @staticmethod
    def draw_rectangle(image, rectangle):
        painter = QPainter()
        painter.begin(image)
        painter.setPen(QPen(QtCore.Qt.red, 3))
        painter.drawRect(*rectangle)
        painter.end()

    def show_check_result_image(self):
        """
        单击单选按钮，对应显示出标志图片
        """
        pairs = self.backend.pairs
        pairs_index = self.backend.pairs_index
        pair = pairs[pairs_index]
        status = pair["status"]

        # 将所有单选按钮置于未选中状态
        self.radioButton_same_person.setAutoExclusive(False)
        self.radioButton_same_person.setChecked(False)
        self.radioButton_same_person.setAutoExclusive(True)
        self.radioButton_different_person.setAutoExclusive(False)
        self.radioButton_different_person.setChecked(False)
        self.radioButton_different_person.setAutoExclusive(False)
        self.radioButton_uncertain.setAutoExclusive(False)
        self.radioButton_uncertain.setChecked(False)
        self.radioButton_uncertain.setAutoExclusive(False)

        if status == "unchecked":
            self.label_result_image.setText("未标注")
        elif status == "right":
            self.radioButton_same_person.setChecked(True)
            self.label_result_image.setPixmap(QPixmap("../images/right.png"))
        elif status == "wrong":
            self.radioButton_different_person.setChecked(True)
            self.label_result_image.setPixmap(QPixmap("../images/wrong.png"))
        elif status == "uncertain":
            self.radioButton_uncertain.setChecked(True)
            self.label_result_image.setPixmap(QPixmap("../images/question.png"))
        else:
            raise ValueError

    def show_image(self):
        pairs = self.backend.pairs
        pairs_index = self.backend.pairs_index
        pair = pairs[pairs_index]
        self.lineEdit_index.setText(F"{pairs_index + 1}/{len(pairs)}")
        self.lineEdit_similarity.setText(pair["distance"])

        gallery = pair["gallery"]
        gallery_images, gallery_rectangles, gallery_index = \
            gallery["image"], gallery["rectangle"], gallery["index"]
        gallery_image = QImage(gallery_images[gallery_index])
        gallery_rectangle = gallery_rectangles[gallery_index]
        self.draw_rectangle(gallery_image, gallery_rectangle)
        self.label_gallery.setPixmap(QPixmap(gallery_image))
        self.lineEdit_gallery_index.setText(F"{gallery_index + 1}/{len(gallery_images)}")
        self.lineEdit_gallery_time.setText(gallery["time"])
        self.lineEdit_gallery_location.setText(gallery["location"])

        query = pair["query"]
        query_images, query_rectangles, query_index = \
            query["image"], query["rectangle"], query["index"]
        query_image = QImage(query_images[query_index])
        query_rectangle = query_rectangles[query_index]
        self.draw_rectangle(query_image, query_rectangle)
        self.label_query.setPixmap(QPixmap(query_image))
        self.lineEdit_query_index.setText(F"{query_index+1}/{len(query_images)}")
        self.lineEdit_query_time.setText(query["time"])
        self.lineEdit_query_location.setText(query["location"])

        self.show_check_result_image()

    def previous_pair_button(self):
        pairs_index = self.backend.pairs_index
        if pairs_index > 0:
            self.backend.pairs_index = pairs_index - 1
        self.show_image()

    def next_pair_button(self):
        pairs = self.backend.pairs
        pairs_index = self.backend.pairs_index

        if pairs[pairs_index]["status"] == "unchecked":  # 必须将当前图片标注后，才能切换进入下一对图片
            return

        if pairs_index < len(pairs) - 1:
            self.backend.pairs_index = pairs_index + 1
        self.show_image()
        self.backend.write_checked_result()

    def previous_gallery_button(self):
        pairs = self.backend.pairs
        pairs_index = self.backend.pairs_index
        pair = pairs[pairs_index]
        if pair["gallery"]["index"] > 0:
            pair["gallery"]["index"] -= 1
        self.show_image()

    def next_gallery_button(self):
        pairs = self.backend.pairs
        pairs_index = self.backend.pairs_index
        pair = pairs[pairs_index]
        if pair["gallery"]["index"] < len(pair["gallery"]["image"]) - 1:
            pair["gallery"]["index"] += 1
        self.show_image()

    def previous_query_button(self):
        pairs = self.backend.pairs
        pairs_index = self.backend.pairs_index
        pair = pairs[pairs_index]
        if pair["query"]["index"] > 0:
            pair["query"]["index"] -= 1
        self.show_image()

    def next_query_button(self):
        pairs = self.backend.pairs
        pairs_index = self.backend.pairs_index
        pair = pairs[pairs_index]
        if pair["query"]["index"] < len(pair["query"]["image"]) - 1:
            pair["query"]["index"] += 1
        self.show_image()

    def same_person_radio_button(self):
        pairs = self.backend.pairs
        pairs_index = self.backend.pairs_index
        pairs[pairs_index]["status"] = "right"
        self.show_check_result_image()

    def different_person_radio_button(self):
        pairs = self.backend.pairs
        pairs_index = self.backend.pairs_index
        pairs[pairs_index]["status"] = "wrong"
        self.show_check_result_image()

    def uncertain_radio_button(self):
        pairs = self.backend.pairs
        pairs_index = self.backend.pairs_index
        pairs[pairs_index]["status"] = "uncertain"
        self.show_check_result_image()

    def closeEvent(self, a0):
        self.backend.write_checked_result()
        super(MatchViewWindow, self).closeEvent(a0)

    def event(self, event):
        if isinstance(event, QtGui.QKeyEvent) and event.type() == QtCore.QEvent.KeyPress:
            if event.key() == Qt.Key_Q:
                self.previous_gallery_button()
            elif event.key() == Qt.Key_E:
                self.next_gallery_button()
            elif event.key() == Qt.Key_A:
                self.previous_query_button()
            elif event.key() == Qt.Key_D:
                self.next_query_button()
            else:
                pass
        return super().event(event)


def main():
    bus_date_folder = sys.argv[1]
    yaml_file_path = os.path.join(bus_date_folder, "information", "label", "od_label.yaml")
    app = QApplication(sys.argv)
    window = MatchViewWindow(yaml_file_path)
    window.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()
